from sklearn import metrics
from sklearn import linear_model
from matplotlib import pyplot as plt
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split as split


def binary_class_auc(X, y, model):
    y_score = model.predict(X)
    print(y_score)
    # y_one_hot = label_binarize(y, np.arange(class_number))
    fpr, tpr, threshold = metrics.roc_curve(y, y_score)
    roc_auc = metrics.auc(fpr, tpr)
    plt.plot(fpr, tpr, label='AUC=%0.3f' % roc_auc)
    plt.plot([0, 1], [0, 1], linestyle='--', lw=2, color='r', alpha=.8)
    plt.legend(loc='lower right')
    plt.show()


if __name__ == '__main__':
    # 1.导入数据集
    iris = load_iris()
    X = iris.data
    y = iris.target
    for i in range(0, len(y)):
        if y[i] == 2:
            y[i] = 1
    # 2.分割训练集
    train_X, test_X, train_y, test_y = split(X, y, test_size=0.2, random_state=0)
    # 3.初始化模型
    model = linear_model.LinearRegression()
    # 4.进行训练
    model.fit(train_X, train_y)
    # 5. 进行预测     # result = model.predict(test_X)
    binary_class_auc(test_X, test_y, model)
